package org.apache.solr.update.processor;

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import org.apache.solr.common.SolrException;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.core.SolrCore;
import org.apache.solr.core.SolrResourceLoader;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.LocalSolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.update.*;
import org.apache.solr.util.plugin.SolrCoreAware;

import org.apache.commons.lang.StringUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.FilenameUtils;

import javax.script.Invocable;
import javax.script.ScriptEngine;
import javax.script.ScriptEngineManager;
import javax.script.ScriptEngineFactory;
import javax.script.ScriptException;

import java.io.IOException;
import java.io.InputStream;
import java.io.Reader;
import java.util.Set;
import java.util.LinkedHashSet;
import java.util.ArrayList;
import java.util.List;
import java.util.Collection;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * <p>
 * An update request processor factory that enables the use of update processors
 * implemented as scripts which can be loaded by the {@link SolrResourceLoader}
 * (usually via the <code>conf</code> dir for the SolrCore).
 * </p>
 * <p>
 * This factory requires at least one configuration parameter named
 * <code>script</code> which may be the name of a script file as a string, or an
 * array of multiple script files. If multiple script files are specified, they
 * are executed sequentially in the order specified in the configuration -- as
 * if multiple factories were configured sequentially
 * </p>
 * <p>
 * Each script file is expected to declare functions with the same name as each
 * method in {@link UpdateRequestProcessor}, using the same arguments. One
 * slight deviation is in the optional return value from these functions: If a
 * script function has a <code>boolean</code> return value, and that value is
 * <code>false</code> then the processor will cleanly terminate processing of
 * the command and return, without forwarding the command on to the next script
 * or processor in the chain. Due to limitations in the {@link ScriptEngine} API
 * used by this factory, it can not enforce that all functions exist on
 * initialization, so errors from missing functions will only be generated at
 * runtime when the chain attempts to use them.
 * </p>
 * <p>
 * The factory may also be configured with an optional "params" argument, which
 * can be an {@link NamedList} (or array, or any other simple Java object) which
 * will be put into the global scope for each script.
 * </p>
 * <p>
 * The following variables are define as global variables for each script:
 * <ul>
 * <li>req - The {@link SolrQueryRequest}</li>
 * <li>rsp - The {@link SolrQueryResponse}</li>
 * <li>logger - A {@link Logger} that can be used for logging purposes in the
 * script</li>
 * <li>params - The "params" init argument in the factory configuration (if any)
 * </li>
 * </ul>
 * </p>
 * <p>
 * Internally this update processor uses JDK 6 scripting engine support, and any
 * {@link Invocable} implementations of <code>ScriptEngine</code> that can be
 * loaded using the Solr Plugin ClassLoader may be used. By default, the engine
 * used for each script is determined by the filed extension (ie: a *.js file
 * will be treated as a JavaScript script) but this can be overridden by
 * specifying an explicit "engine" name init param for the factory, which
 * identifies a registered name of a {@link ScriptEngineFactory}. (This may be
 * particularly useful if multiple engines are available for the same scripting
 * language, and you wish to force the usage of a particular engine because of
 * known quirks)
 * </p>
 * <p>
 * A new {@link ScriptEngineManager} is created for each
 * <code>SolrQueryRequest</code> defining a "global" scope for the script(s)
 * which is request specific. Separate <code>ScriptEngine</code> instances are
 * then used to evaluate the script files, resulting in an "engine" scope that
 * is specific to each script.
 * </p>
 * <p>
 * A simple example...
 * </p>
 * 
 * <pre class="prettyprint">
 * &lt;processor class="solr.StatelessScriptUpdateProcessorFactory"&gt;
 *   &lt;str name="script"&gt;updateProcessor.js&lt;/str&gt;
 * &lt;/processor&gt;
 * </pre>
 * <p>
 * A more complex example involving multiple scripts in different languages, and
 * a "params" <code>NamedList</code> that will be put into the global scope of
 * each script...
 * </p>
 * 
 * <pre class="prettyprint">
 * &lt;processor class="solr.StatelessScriptUpdateProcessorFactory"&gt;
 *   &lt;arr name="script"&gt;
 *     &lt;str name="script"&gt;first-processor.js&lt;/str&gt;
 *     &lt;str name="script"&gt;second-processor.py&lt;/str&gt;
 *   &lt;/arr&gt;
 *   &lt;lst name="params"&gt;
 *     &lt;bool name="a_bool_value"&gt;true&lt;/bool&gt;
 *     &lt;int name="and_int_value"&gt;3&lt;/int&gt;
 *   &lt;/lst&gt;
 * &lt;/processor&gt;
 * </pre>
 * <p>
 * An example where the script file extensions are ignored, and an explicit
 * script engine is used....
 * </p>
 * 
 * <pre class="prettyprint">
 * &lt;processor class="solr.StatelessScriptUpdateProcessorFactory"&gt;
 *   &lt;arr name="script"&gt;
 *     &lt;str name="script"&gt;first-processor.txt&lt;/str&gt;
 *     &lt;str name="script"&gt;second-processor.txt&lt;/str&gt;
 *   &lt;/arr&gt;
 *   &lt;str name="engine"&gt;rhino&lt;/str&gt;
 * &lt;/processor&gt;
 * </pre>
 * 
 */
public class StatelessScriptUpdateProcessorFactory extends
		UpdateRequestProcessorFactory implements SolrCoreAware {

	public static Logger log = LoggerFactory
			.getLogger(StatelessScriptUpdateProcessorFactory.class);

	private final static String SCRIPT_ARG = "script";
	private final static String PARAMS_ARG = "params";
	private final static String ENGINE_NAME_ARG = "engine";

	private List<ScriptFile> scriptFiles;

	/** if non null, this is an override for the engine for all scripts */
	private String engineName = null;

	private Object params = null;

	private SolrResourceLoader resourceLoader;

	private ScriptEngineCustomizer scriptEngineCustomizer;

	@Override
	public void init(NamedList args) {
		Collection<String> scripts = FieldMutatingUpdateProcessorFactory
				.oneOrMany(args, SCRIPT_ARG);
		if (scripts.isEmpty()) {
			throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
					"StatelessScriptUpdateProcessorFactory must be "
							+ "initialized with at least one " + SCRIPT_ARG);
		}
		scriptFiles = new ArrayList<ScriptFile>();
		for (String script : scripts) {
			scriptFiles.add(new ScriptFile(script));
		}

		params = args.remove(PARAMS_ARG);

		Object engine = args.remove(ENGINE_NAME_ARG);
		if (engine != null) {
			if (engine instanceof String) {
				engineName = (String) engine;
			} else {
				throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
						"'" + ENGINE_NAME_ARG
								+ "' init param must be a String (found: "
								+ engine.getClass() + ")");
			}
		}

		super.init(args);

	}

	@Override
	public UpdateRequestProcessor getInstance(SolrQueryRequest req,
			SolrQueryResponse rsp, UpdateRequestProcessor next) {
		List<EngineInfo> scriptEngines = null;

		scriptEngines = initEngines(req, rsp);

		return new ScriptUpdateProcessor(req, rsp, scriptEngines, next);
	}

	// TODO: Make this useful outside of tests, such that a
	// ScriptEngineCustomizer could be looked up through the resource loader
	void setScriptEngineCustomizer(ScriptEngineCustomizer scriptEngineCustomizer) {
		this.scriptEngineCustomizer = scriptEngineCustomizer;
	}

	@Override
	public void inform(SolrCore core) {
		resourceLoader = core.getResourceLoader();

		// test that our engines & scripts are valid

		SolrQueryResponse rsp = new SolrQueryResponse();
		SolrQueryRequest req = new LocalSolrQueryRequest(core,
				new ModifiableSolrParams());
		try {
			initEngines(req, rsp);
		} catch (Exception e) {
			String msg = "Unable to initialize scripts: " + e.getMessage();
			log.error(msg, e);
			throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, msg,
					e);
		} finally {
			req.close();
		}

	}

	// ================================================ Helper Methods
	// ==================================================

	/**
	 * Initializes a list of script engines - an engine per script file.
	 * 
	 * @param req
	 *            The solr request.
	 * @param rsp
	 *            The solr response
	 * @return The list of initialized script engines.
	 */
	private List<EngineInfo> initEngines(SolrQueryRequest req,
			SolrQueryResponse rsp) throws SolrException {

		List<EngineInfo> scriptEngines = new ArrayList<EngineInfo>();

		ScriptEngineManager scriptEngineManager = new ScriptEngineManager(
				resourceLoader.getClassLoader());

		scriptEngineManager.put("logger", log);
		scriptEngineManager.put("req", req);
		scriptEngineManager.put("rsp", rsp);
		if (params != null) {
			scriptEngineManager.put("params", params);
		}

		for (ScriptFile scriptFile : scriptFiles) {
			ScriptEngine engine = null;
			if (null != engineName) {
				engine = scriptEngineManager.getEngineByName(engineName);
				if (engine == null) {
					String details = getSupportedEngines(scriptEngineManager,
							false);
					throw new SolrException(
							SolrException.ErrorCode.SERVER_ERROR,
							"No ScriptEngine found by name: "
									+ engineName
									+ (null != details ? " -- supported names: "
											+ details
											: ""));
				}
			} else {
				engine = scriptEngineManager.getEngineByExtension(scriptFile
						.getExtension());
				if (engine == null) {
					String details = getSupportedEngines(scriptEngineManager,
							true);
					throw new SolrException(
							SolrException.ErrorCode.SERVER_ERROR,
							"No ScriptEngine found by file extension: "
									+ scriptFile.getFileName()
									+ (null != details ? " -- supported extensions: "
											+ details
											: ""));

				}
			}

			if (!(engine instanceof Invocable)) {
				String msg = "Engine "
						+ ((null != engineName) ? engineName
								: ("for script " + scriptFile.getFileName()))
						+ " does not support function invocation (via Invocable): "
						+ engine.getClass().toString() + " ("
						+ engine.getFactory().getEngineName() + ")";
				log.error(msg);
				throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
						msg);
			}

			if (scriptEngineCustomizer != null) {
				scriptEngineCustomizer.customize(engine);
			}

			scriptEngines.add(new EngineInfo((Invocable) engine, scriptFile));
			try {
				Reader scriptSrc = scriptFile.openReader(resourceLoader);

				try {
					engine.eval(scriptSrc);
				} catch (ScriptException e) {
					throw new SolrException(
							SolrException.ErrorCode.SERVER_ERROR,
							"Unable to evaluate script: "
									+ scriptFile.getFileName(), e);
				} finally {
					IOUtils.closeQuietly(scriptSrc);
				}
			} catch (IOException ioe) {
				throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
						"Unable to evaluate script: "
								+ scriptFile.getFileName(), ioe);
			}
		}
		return scriptEngines;
	}

	/**
	 * For error messages - returns null if there are any exceptions of any kind
	 * building the string (or of the list is empty for some unknown reason).
	 * 
	 * @param ext
	 *            - if true, list of extensions, otherwise a list of engine
	 *            names
	 */
	private static String getSupportedEngines(ScriptEngineManager mgr,
			boolean ext) {
		String result = null;
		try {
			List<ScriptEngineFactory> factories = mgr.getEngineFactories();
			if (null == factories)
				return result;

			Set<String> engines = new LinkedHashSet<String>(factories.size());
			for (ScriptEngineFactory f : factories) {
				if (ext) {
					engines.addAll(f.getExtensions());
				} else {
					engines.addAll(f.getNames());
				}
			}
			result = StringUtils.join(engines, ", ");
		} catch (RuntimeException e) {
			/* :NOOP: */
		}
		return result;
	}

	// ================================================= Inner Classes
	// ==================================================

	/**
	 * The actual update processor. All methods delegate to scripts.
	 */
	private static class ScriptUpdateProcessor extends UpdateRequestProcessor {

		private List<EngineInfo> engines;

		private ScriptUpdateProcessor(SolrQueryRequest req,
				SolrQueryResponse res, List<EngineInfo> engines,
				UpdateRequestProcessor next) {
			super(next);
			this.engines = engines;
		}

		@Override
		public void processAdd(AddUpdateCommand cmd) throws IOException {
			if (invokeFunction("processAdd", cmd)) {
				super.processAdd(cmd);
			}
		}

		@Override
		public void processDelete(DeleteUpdateCommand cmd) throws IOException {
			if (invokeFunction("processDelete", cmd)) {
				super.processDelete(cmd);
			}

		}

		@Override
		public void processMergeIndexes(MergeIndexesCommand cmd)
				throws IOException {
			if (invokeFunction("processMergeIndexes", cmd)) {
				super.processMergeIndexes(cmd);
			}
		}

		@Override
		public void processCommit(CommitUpdateCommand cmd) throws IOException {
			if (invokeFunction("processCommit", cmd)) {
				super.processCommit(cmd);
			}
		}

		@Override
		public void processRollback(RollbackUpdateCommand cmd)
				throws IOException {
			if (invokeFunction("processRollback", cmd)) {
				super.processRollback(cmd);
			}
		}

		@Override
		public void finish() throws IOException {
			if (invokeFunction("finish")) {
				super.finish();
			}
		}

		/**
		 * returns true if processing should continue, or false if the request
		 * should be ended now. Result value is computed from the return value
		 * of the script function if: it exists, is non-null, and can be cast to
		 * a java Boolean.
		 */
		private boolean invokeFunction(String name, Object... cmd) {

			for (EngineInfo engine : engines) {
				try {
					Object result = engine.getEngine()
							.invokeFunction(name, cmd);
					if (null != result && result instanceof Boolean) {
						if (!((Boolean) result).booleanValue()) {
							return false;
						}
					}

				} catch (ScriptException e) {
					throw new SolrException(
							SolrException.ErrorCode.SERVER_ERROR,
							"Unable to invoke function " + name
									+ " in script: "
									+ engine.getScriptFile().getFileName()
									+ ": " + e.getMessage(), e);
				} catch (NoSuchMethodException e) {
					throw new SolrException(
							SolrException.ErrorCode.SERVER_ERROR,
							"Unable to invoke function " + name
									+ " in script: "
									+ engine.getScriptFile().getFileName()
									+ ": " + e.getMessage(), e);
				}
			}

			return true;
		}
	}

	/**
	 * Holds the script engine and its associated script file.
	 */
	private static class EngineInfo {

		private final Invocable engine;
		private final ScriptFile scriptFile;

		private EngineInfo(Invocable engine, ScriptFile scriptFile) {
			this.engine = engine;
			this.scriptFile = scriptFile;
		}

		public Invocable getEngine() {
			return engine;
		}

		public ScriptFile getScriptFile() {
			return scriptFile;
		}
	}

	/**
	 * Represents a script file.
	 */
	private static class ScriptFile {

		private final String fileName;
		private final String extension;

		private ScriptFile(String fileName) {
			this.fileName = fileName;
			this.extension = FilenameUtils.getExtension(fileName);
		}

		public String getFileName() {
			return fileName;
		}

		public String getExtension() {
			return extension;
		}

		public Reader openReader(SolrResourceLoader resourceLoader)
				throws IOException {
			InputStream input = resourceLoader.openResource(fileName);
			return org.apache.lucene.util.IOUtils.getDecodingReader(input,
					org.apache.lucene.util.IOUtils.CHARSET_UTF_8);
		}
	}
}
